package com.beyondsoft.sabg.edu.utils;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.CompletableFuture;

/**
 * SSE流式HTTP客户端工具类
 */
@Slf4j
public class SseHttpClient {

    /**
     * 发送SSE流式请求
     * @param url 请求URL
     * @param requestBody 请求体JSON
     * @param apiKey API密钥
     * @param emitter SSE发射器
     */
    public static void sendSseRequest(String url, String requestBody, String apiKey, SseEmitter emitter) {
        CompletableFuture.runAsync(() -> {
            HttpURLConnection connection = null;
            BufferedReader reader = null;
            
            try {
                // 创建连接
                URL urlObj = new URL(url);
                connection = (HttpURLConnection) urlObj.openConnection();
                
                // 设置请求方法和头部
                connection.setRequestMethod("POST");
                connection.setRequestProperty("Content-Type", "application/json");
                connection.setRequestProperty("Authorization", "Bearer " + apiKey);
                connection.setRequestProperty("Accept", "text/event-stream");
                connection.setRequestProperty("Cache-Control", "no-cache");
                connection.setDoOutput(true);
                connection.setDoInput(true);
                connection.setConnectTimeout(30000); // 30秒连接超时
                connection.setReadTimeout(300000);   // 5分钟读取超时
                
                // 发送请求体
                try (OutputStream os = connection.getOutputStream()) {
                    byte[] input = requestBody.getBytes(StandardCharsets.UTF_8);
                    os.write(input, 0, input.length);
                }
                
                // 检查响应码
                int responseCode = connection.getResponseCode();
                log.info("LLM API响应码: {}", responseCode);
                
                if (responseCode != HttpURLConnection.HTTP_OK) {
                    String errorMessage = "LLM API请求失败，响应码: " + responseCode;
                    log.error(errorMessage);
                    emitter.completeWithError(new RuntimeException(errorMessage));
                    return;
                }
                
                // 读取流式响应
                reader = new BufferedReader(new InputStreamReader(connection.getInputStream(), StandardCharsets.UTF_8));
                String line;
                
                while ((line = reader.readLine()) != null) {
                    if (StrUtil.isNotBlank(line)) {
                        // 处理SSE数据
                        String processedData = processSseData(line);
                        if (StrUtil.isNotBlank(processedData)) {
                            try {
                                emitter.send(SseEmitter.event().data(processedData));
                            } catch (IOException e) {
                                log.error("发送SSE数据失败", e);
                                break;
                            }
                        }
                    }
                }
                
                // 完成流式响应
                emitter.complete();
                log.info("SSE流式响应处理完成");
                
            } catch (Exception e) {
                log.error("SSE请求处理失败", e);
                try {
                    emitter.completeWithError(e);
                } catch (Exception ex) {
                    log.error("完成SSE连接失败", ex);
                }
            } finally {
                // 关闭资源
                if (reader != null) {
                    try {
                        reader.close();
                    } catch (IOException e) {
                        log.error("关闭reader失败", e);
                    }
                }
                if (connection != null) {
                    connection.disconnect();
                }
            }
        });
    }
    
    /**
     * 处理SSE数据格式
     * @param rawData 原始数据
     * @return 处理后的数据
     */
    private static String processSseData(String rawData) {
        if (StrUtil.isBlank(rawData)) {
            return "";
        }
        
        try {
            // 处理SSE格式的数据
            if (rawData.startsWith("data: ")) {
                String jsonData = rawData.substring(6).trim();
                
                // 检查是否是结束标记
                if ("[DONE]".equals(jsonData) || "null".equals(jsonData)) {
                    return "";
                }
                
                if (StrUtil.isNotBlank(jsonData)) {
                    try {
                        JSONObject jsonObject = new JSONObject(jsonData);
                        
                        // 根据不同的事件类型处理数据
                        String event = jsonObject.getStr("event");
                        if ("message".equals(event)) {
                            return jsonObject.getStr("answer", "");
                        } else if ("message_end".equals(event)) {
                            // 消息结束，返回空字符串
                            return "";
                        } else if ("error".equals(event)) {
                            // 错误事件
                            log.error("LLM API返回错误: {}", jsonData);
                            return "";
                        }
                    } catch (Exception e) {
                        // 如果不是JSON格式，可能是纯文本数据
                        log.debug("非JSON格式数据: {}", jsonData);
                        return jsonData;
                    }
                }
            } else if (rawData.startsWith("event: ") || rawData.startsWith("id: ") || rawData.startsWith("retry: ")) {
                // SSE元数据，忽略
                return "";
            } else if (StrUtil.isNotBlank(rawData.trim())) {
                // 其他非空数据，直接返回
                return rawData;
            }
        } catch (Exception e) {
            log.warn("解析SSE数据失败: {}", rawData, e);
            // 如果解析失败，直接返回原始数据
            return rawData;
        }
        
        return "";
    }
} 